/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file fused_add_topk_div_tiling.cpp
 * \brief
 */
#include "fused_add_topk_div_tiling.h"

namespace {
const static uint32_t BYTE_BLOCK = 32;
const static int32_t X_INPUT_INDEX = 0;
const static int32_t ADD_NUM_INPUT_INDEX = 1;
const static int32_t MAPPING_NUM_INPUT_IDX = 2;
const static int32_t MAPPING_TABLE_INPUT_IDX = 3;
const static uint32_t DIM_INDEX0 = 0;
const static uint32_t DIM_INDEX1 = 1;
const static uint32_t TWO_DIM = 2;

const static int32_t FLOAT_BYTES = 4;
const static uint32_t RESERVED_UB = 16 * 1024;
const static int32_t BASE_COUNT = 256;
const static int32_t NUM_FOUR = 4;
const static int32_t SORT_UNIT = 32;
const static int32_t GROUP_LIMIT = 8;
const static int32_t MAPPING_TABLE_DIM_ONE_LIMIT = 128;
const static int32_t NUM_SIXTEEN = 16;
const static int32_t NUM_TEN = 10;
const static uint32_t SYS_WORKSPACESIZE = 16 * 1024 * 1024;

const static size_t ATTR_GROUP_NUM_IDX = 0;
const static size_t ATTR_GROUP_TOPK_IDX = 1;
const static size_t ATTR_N_IDX = 2;
const static size_t ATTR_K_IDX = 3;
const static size_t ATTR_ACTIVATE_TYPE_IDX = 4;
const static size_t ATTR_IS_NORM_IDX = 5;
const static size_t ATTR_ENABLE_EXPERT_MAPPING_IDX = 6;
const static size_t ATTR_SCALE_IDX = 7;

static std::map<ge::DataType, int32_t> g_dtypeMap = {{ge::DT_FLOAT, 0},
                                                     {ge::DT_FLOAT16, 1},
                                                     {ge::DT_BF16, 2}};
} // namespace

namespace optiling {
class FusedAddTopkDivTiling {
public:
    explicit FusedAddTopkDivTiling(gert::TilingContext* context) : tilingcontext(context){};
    ge::graphStatus Init();
    ge::graphStatus RunKernelTiling();
    void PrintTilingData() const;
private:
    FusedAddTopkDivTilingData tilingData;
    gert::TilingContext* tilingcontext = nullptr;
    ge::graphStatus GetAndCheckAttr();
    ge::graphStatus CheckAttr();
    void GetTilingKey();
    void GetUsedCore();
    void SplitUb();
    void FillTilingData();
    template <typename T1, typename T2>
    inline T1 CeilAlign(T1 a, T2 b) const
    {
        return b == 0 ? a : (a + b - 1) / b * b;
    }
    template <typename T1, typename T2>
    inline T1 FloorAlign(T1 a, T2 b) const
    {
        return b == 0 ? a : a / b * b;
    }

    uint32_t firstDimSize = 0;
    uint32_t secondDimSize = 0;
    uint32_t addNumDimSize = 0;
    uint32_t expertNum = 0;
    uint32_t tableDim = 0;
    uint32_t groupNum = 1;
    uint32_t groupTopk = 1;
    uint32_t n = 1;
    uint32_t k = 0;
    uint32_t activateType = 0;
    uint32_t isNorm = 1;
    float scale = 1.0;
    uint32_t groupEles = 0;
    uint32_t enableExpertMapping = 0;
    uint32_t ubSize = 0;
    uint32_t usedCoreNum = 0;
    uint32_t coreNum = 0;
    uint32_t batchPerCore = 1;
    uint32_t tailBatch = 0;
    uint32_t ubFactorElement = 0;
    uint32_t tilingKey = 0;
    uint32_t dtype = 1;
    uint64_t workspacePerCore = 0;
    const char *opName = nullptr;
};

ge::graphStatus FusedAddTopkDivTiling::Init() {
    opName = tilingcontext->GetNodeName();
    OPS_LOG_D(opName, "TilingForFusedAddTopkDiv init.");
    auto platformInfo = platform_ascendc::PlatformAscendC(tilingcontext->GetPlatformInfo());
    coreNum = platformInfo.GetCoreNumAiv();
    uint64_t platformUbSize = 0;
    platformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, platformUbSize);
    uint32_t avaliableUb = static_cast<uint32_t>(platformUbSize) - RESERVED_UB;
    ubSize = FloorAlign(avaliableUb, BYTE_BLOCK);

    auto xShape = tilingcontext->GetInputShape(X_INPUT_INDEX)->GetStorageShape();
    if (xShape.GetDimNum() != TWO_DIM) {
        OPS_LOG_E(opName, "the dimNum of x should be 2, but got %ld.", xShape.GetDimNum());
        return ge::GRAPH_FAILED;
    }
    auto addNumShape = tilingcontext->GetInputShape(ADD_NUM_INPUT_INDEX)->GetStorageShape();
    if (addNumShape.GetDimNum() != 1) {
        OPS_LOG_E(opName, "the dimNum of add_num should be 1, but got %ld.", addNumShape.GetDimNum());
        return ge::GRAPH_FAILED;
    }
    firstDimSize = xShape.GetDim(DIM_INDEX0);
    secondDimSize = xShape.GetDim(DIM_INDEX1);
    addNumDimSize = addNumShape.GetDim(DIM_INDEX0);
    if (secondDimSize != addNumDimSize) {
        OPS_LOG_E(opName, "add_num.shape[1] should be equal to x.shape[1].");
        return ge::GRAPH_FAILED;
    }
    dtype = g_dtypeMap[tilingcontext->GetInputDesc(0)->GetDataType()];

    return ge::GRAPH_SUCCESS;
}

ge::graphStatus FusedAddTopkDivTiling::GetAndCheckAttr() {
    OPS_LOG_D(opName, "TilingForFusedAddTopkDiv GetAndCheckAttr.");
    const gert::RuntimeAttrs *attrs = tilingcontext->GetAttrs();
    OPS_LOG_E_IF_NULL(tilingcontext, attrs, return ge::GRAPH_FAILED);
    groupNum = static_cast<uint32_t>(*(attrs->GetAttrPointer<int64_t>(ATTR_GROUP_NUM_IDX)));
    groupTopk = static_cast<uint32_t>(*(attrs->GetAttrPointer<int64_t>(ATTR_GROUP_TOPK_IDX)));
    n = static_cast<uint32_t>(*(attrs->GetAttrPointer<int64_t>(ATTR_N_IDX)));
    k = static_cast<uint32_t>(*(attrs->GetAttrPointer<int64_t>(ATTR_K_IDX)));
    activateType = static_cast<uint32_t>(*(attrs->GetAttrPointer<int64_t>(ATTR_ACTIVATE_TYPE_IDX)));
    isNorm = static_cast<uint32_t>(*(attrs->GetAttrPointer<bool>(ATTR_IS_NORM_IDX)));
    enableExpertMapping = static_cast<uint32_t>(*(attrs->GetAttrPointer<bool>(ATTR_ENABLE_EXPERT_MAPPING_IDX)));
    groupEles = groupNum == 0 ? secondDimSize : secondDimSize / groupNum;
    scale = *(attrs->GetAttrPointer<float>(ATTR_SCALE_IDX));
    
    if (enableExpertMapping) {
        auto mappingNumShape = tilingcontext->GetInputShape(MAPPING_NUM_INPUT_IDX)->GetStorageShape();
        if (mappingNumShape.GetDimNum() != 1) {
            OPS_LOG_E(opName, "the dimNum of mapping_num should be 1, but got %ld.", mappingNumShape.GetDimNum());
            return ge::GRAPH_FAILED;
        }
        if (mappingNumShape.GetDim(DIM_INDEX0) != secondDimSize) {
            OPS_LOG_E(opName, "mapping_num.shape[0] should be equal to x.shape[1].");
            return ge::GRAPH_FAILED;
        }
        auto mappingTableShape = tilingcontext->GetInputShape(MAPPING_TABLE_INPUT_IDX)->GetStorageShape();
        if (mappingTableShape.GetDimNum() != TWO_DIM) {
            OPS_LOG_E(opName, "the dimNum of mapping_table should be 2, but got %ld.", mappingTableShape.GetDimNum());
            return ge::GRAPH_FAILED;
        }
        expertNum = mappingTableShape.GetDim(DIM_INDEX0);
        tableDim = mappingTableShape.GetDim(DIM_INDEX1);
        if (expertNum != secondDimSize) {
            OPS_LOG_E(opName, "mapping_table.shape[0] should be equal to x.shape[1].");
            return ge::GRAPH_FAILED;
        }
        if (tableDim > MAPPING_TABLE_DIM_ONE_LIMIT) {
            OPS_LOG_E(opName, "mapping_table.shape[1] should less or equal to 128.");
            return ge::GRAPH_FAILED;
        }
    }
    return CheckAttr();
}

ge::graphStatus FusedAddTopkDivTiling::CheckAttr() {
    if (groupNum <= 0) {
        OPS_LOG_E(opName, "groupNum should be greater than 0, but got %u.", groupNum);
        return ge::GRAPH_FAILED;
    }
    if (groupTopk <= 0) {
        OPS_LOG_E(opName, "groupTopk should be greater than 0, but got %u.", groupTopk);
        return ge::GRAPH_FAILED;
    }
    if (n <= 0) {
        OPS_LOG_E(opName, "n should be greater than 0, but got %u.", n);
        return ge::GRAPH_FAILED;
    }
    if (k <= 0) {
        OPS_LOG_E(opName, "k should be greater than 0, but got %u.", k);
        return ge::GRAPH_FAILED;
    }
    if (secondDimSize % groupNum != 0) {
        OPS_LOG_E(opName, "x.shape[1] should be a multiple of groupNum.");
        return ge::GRAPH_FAILED;
    }
    if (groupTopk > groupNum) {
        OPS_LOG_E(opName, "groupTopk should be less or equal to groupNum, but got groupTopk: %u, groupNum: %u.",
                  groupTopk, groupNum);
        return ge::GRAPH_FAILED;
    }
    if (k > secondDimSize) {
        OPS_LOG_E(opName, "k should be less or equal to x.shape[1], but got k: %u, x.shape[1]: %u.",
                  k, secondDimSize);
        return ge::GRAPH_FAILED;
    }
    if (n > secondDimSize / groupNum) {
        OPS_LOG_E(opName,
                  "n should be less or equal to x.shape[1] / groupNum, but got n: %u, secondDimSize / groupNum: %u.",
                  n, secondDimSize / groupNum);
        return ge::GRAPH_FAILED;
    }
    if (secondDimSize / groupNum > SORT_UNIT) {
        OPS_LOG_E(opName, "x.shape[1] / groupNum should be less or equal to 32, but got %u.", secondDimSize / groupNum);
        return ge::GRAPH_FAILED;
    }
    if (activateType != 0) {
        OPS_LOG_E(opName, "activateType currently only suppoorts equal to 0.");
        return ge::GRAPH_FAILED;
    }
    if ((secondDimSize >= SORT_UNIT) && (groupNum != GROUP_LIMIT)) {
        OPS_LOG_E(opName, "when x.shape[1] is greater or equal to 32, groupNum only supports 8.");
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}

void FusedAddTopkDivTiling::GetTilingKey() {
    tilingKey = enableExpertMapping * NUM_TEN + dtype;
}

void FusedAddTopkDivTiling::GetUsedCore() {
    if (firstDimSize <= coreNum) {
        batchPerCore = 1;
        usedCoreNum = firstDimSize;
        tailBatch = 0;
        return;
    }
    batchPerCore = coreNum == 0 ? firstDimSize : firstDimSize / coreNum;
    tailBatch = firstDimSize % coreNum;
    usedCoreNum = coreNum;
}

void FusedAddTopkDivTiling::SplitUb() {
    uint32_t tilingDataSize = CeilAlign(sizeof(FusedAddTopkDivTilingData), BYTE_BLOCK);
    uint32_t canUseUbSize = FloorAlign(ubSize - tilingDataSize, BYTE_BLOCK);
    ubFactorElement = FloorAlign(canUseUbSize, BYTE_BLOCK);
}

void FusedAddTopkDivTiling::FillTilingData() {
    tilingData.set_firstDimSize(firstDimSize);
    tilingData.set_secondDimSize(secondDimSize);
    tilingData.set_addNumDimSize(addNumDimSize);
    tilingData.set_dtype(dtype);
    tilingData.set_groupNum(groupNum);
    tilingData.set_groupTopk(groupTopk);
    tilingData.set_n(n);
    tilingData.set_k(k);
    tilingData.set_activateType(activateType);
    tilingData.set_isNorm(isNorm);
    tilingData.set_enableExpertMapping(enableExpertMapping);
    tilingData.set_scale(scale);
    tilingData.set_groupEles(groupEles);
    tilingData.set_blockNum(usedCoreNum);
    tilingData.set_ubFactorElement(ubFactorElement);
    tilingData.set_batchPerCore(batchPerCore);
    tilingData.set_tailBatch(tailBatch);
    tilingData.set_tilingKey(tilingKey);
    workspacePerCore = BASE_COUNT * FLOAT_BYTES;
    tilingData.set_workspacePerCore(workspacePerCore);
    tilingData.set_tempSize(firstDimSize * secondDimSize * FLOAT_BYTES);
    tilingData.set_expertNum(expertNum);
    tilingData.set_tableDim(tableDim);
}

void FusedAddTopkDivTiling::PrintTilingData() const {
    OPS_LOG_D(opName, "firstDimSize is: %u", firstDimSize);
    OPS_LOG_D(opName, "secondDimSize is: %u", secondDimSize);
    OPS_LOG_D(opName, "addNumDimSize is: %u", addNumDimSize);
    OPS_LOG_D(opName, "groupNum is: %u", groupNum);
    OPS_LOG_D(opName, "groupTopk is: %u", groupTopk);
    OPS_LOG_D(opName, "n is: %u", n);
    OPS_LOG_D(opName, "k is: %u", k);
    OPS_LOG_D(opName, "activateType is: %u", activateType);
    OPS_LOG_D(opName, "isNorm is: %u", isNorm);
    OPS_LOG_D(opName, "scale is: %f", scale);
    OPS_LOG_D(opName, "groupEles is: %u", groupEles);
    OPS_LOG_D(opName, "blockNum is: %u", usedCoreNum);
    OPS_LOG_D(opName, "ubFactorElement is: %u", ubFactorElement);
    OPS_LOG_D(opName, "batchPerCore is: %u", batchPerCore);
    OPS_LOG_D(opName, "tailBatch is: %u", tailBatch);
    OPS_LOG_D(opName, "tilingKey is: %u", tilingKey);
    OPS_LOG_D(opName, "dtype is: %u", dtype);
    OPS_LOG_D(opName, "enableExpertMapping is: %u", enableExpertMapping);
    OPS_LOG_D(opName, "expertNum is: %u", expertNum);
    OPS_LOG_D(opName, "tableDim is: %u", tableDim);
    OPS_LOG_D(opName, "workspacePerCore is: %lu", workspacePerCore);
}

ge::graphStatus FusedAddTopkDivTiling::RunKernelTiling() {
    OPS_LOG_D(opName, "TilingForFusedAddTopkDiv start.");
    if (GetAndCheckAttr() == ge::GRAPH_FAILED) {
        return ge::GRAPH_FAILED;
    }
    GetTilingKey();
    GetUsedCore();
    SplitUb();
    FillTilingData();
    PrintTilingData();
    uint32_t syncWorkspaceSize = SYS_WORKSPACESIZE + usedCoreNum * workspacePerCore;
    size_t* currentWorkspace = tilingcontext->GetWorkspaceSizes(1);
    currentWorkspace[0] = syncWorkspaceSize;

    tilingData.SaveToBuffer(tilingcontext->GetRawTilingData()->GetData(),
                            tilingcontext->GetRawTilingData()->GetCapacity());
    tilingcontext->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
    tilingcontext->SetBlockDim(usedCoreNum);
    tilingcontext->SetTilingKey(tilingData.get_tilingKey());
    return ge::GRAPH_SUCCESS;
}

ASCENDC_EXTERN_C ge::graphStatus TilingForFusedAddTopkDiv(gert::TilingContext* context) {
    FusedAddTopkDivTiling tilingObject(context);
    auto ret = tilingObject.Init();
    if (ret != ge::GRAPH_SUCCESS) {
        OPS_LOG_E(context->GetNodeName(), "tiling Init failed.");
        return ge::GRAPH_FAILED;
    }
    return tilingObject.RunKernelTiling();
}

ASCENDC_EXTERN_C ge::graphStatus TilingPrepare4FusedAddTopkDiv(gert::TilingParseContext* context) {
    return ge::GRAPH_SUCCESS;
}

IMPL_OP_OPTILING(FusedAddTopkDiv)
    .Tiling(TilingForFusedAddTopkDiv)
    .TilingParse<FusedAddTopkDivCompileInfo>(TilingPrepare4FusedAddTopkDiv);
}  // namespace optiling